#!/usr/bin/env python3
"""Run the 2D Helmholtz locELM reproduction experiment."""
from __future__ import annotations

import argparse
from time import perf_counter

import numpy as np

from reproduct_paper import LocELMSolver2D
from reproduct_paper.problems import helmholtz_2d_reference


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--subdomains-x", type=int, default=2, help="number of sub-domains in x")
    parser.add_argument("--subdomains-y", type=int, default=2, help="number of sub-domains in y")
    parser.add_argument("--collocation-x", type=int, default=25, help="collocation points per sub-domain in x")
    parser.add_argument("--collocation-y", type=int, default=25, help="collocation points per sub-domain in y")
    parser.add_argument("--width", type=int, default=400, help="hidden units (training parameters) per sub-domain")
    parser.add_argument("--rand-mag", type=float, default=1.5, help="maximum magnitude of random hidden parameters")
    parser.add_argument("--seed", type=int, default=1, help="seed for RNG")
    parser.add_argument("--samples", type=int, default=200, help="evaluation resolution per direction")
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    config = helmholtz_2d_reference()
    rng = np.random.default_rng(args.seed)

    solver = LocELMSolver2D(
        domain=config["domain"],
        num_subdomains=(args.subdomains_x, args.subdomains_y),
        collocation_per_subdomain=(args.collocation_x, args.collocation_y),
        hidden_width=args.width,
        rand_magnitude=args.rand_mag,
        rhs=lambda x, y: np.asarray(config["rhs"](x, y)).item(),
        boundary=lambda x, y: np.asarray(config["boundary"](x, y)).item(),
        lam=config["lambda"],
        rng=rng,
    )

    start = perf_counter()
    solver.solve()
    elapsed = perf_counter() - start

    # Evaluation grid
    x_eval = np.linspace(config["domain"][0][0], config["domain"][0][1], args.samples)
    y_eval = np.linspace(config["domain"][1][0], config["domain"][1][1], args.samples)
    grid_x, grid_y = np.meshgrid(x_eval, y_eval, indexing="ij")
    numerical = solver.evaluate(grid_x.ravel(), grid_y.ravel())
    exact = config["exact"](grid_x.ravel(), grid_y.ravel())
    error = numerical - exact
    max_error = np.max(np.abs(error))
    rms_error = np.sqrt(np.mean(error ** 2))

    print("locELM 2D Helmholtz experiment")
    print(f"  sub-domains (x,y):    {args.subdomains_x} x {args.subdomains_y}")
    print(f"  collocation/subdom:   {args.collocation_x} x {args.collocation_y}")
    print(f"  hidden width/subdom:  {args.width}")
    print(f"  random magnitude:     {args.rand_mag}")
    print(f"  solve time (s):       {elapsed:.3f}")
    print(f"  max abs error:        {max_error:.3e}")
    print(f"  RMS error:            {rms_error:.3e}")


if __name__ == "__main__":
    main()
